package org.luxor.common.netty.protocol.codec;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.MessageToMessageEncoder;
import org.luxor.common.netty.protocol.*;
import org.luxor.common.netty.protocol.message.NettyConnectMessage;
import org.luxor.common.netty.protocol.message.NettyMessage;
import org.luxor.common.netty.protocol.message.NettyPublishMessage;
import org.luxor.common.netty.utils.ProtostuffSerializeUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;

import static io.netty.buffer.ByteBufUtil.utf8Bytes;
import static org.luxor.common.netty.protocol.NettyCodecUtil.*;

/**
 * NettyMessage 编码器
 *
 * @author Mr.yan @date 2024/1/7
 * @see io.netty.handler.codec.mqtt.MqttEncoder 参考Mqtt消息编码器
 * 参考 http://docs.oasis-open.org/mqtt/mqtt/v5.0/csprd02/mqtt-v5.0-csprd02.html#_Toc498345306
 */
public class NettyEncoder extends MessageToMessageEncoder<NettyMessage> {
    private static final Logger log = LoggerFactory.getLogger(NettyEncoder.class);

    @Override
    protected void encode(ChannelHandlerContext ctx, NettyMessage msg, List<Object> out) throws Exception {
        out.add(doEncode(ctx, msg));
    }

    static ByteBuf doEncode(ChannelHandlerContext ctx, NettyMessage message) {
        switch (message.fixedHeader().messageType()) {
            case CONNECT:
                return encodeConnectMessage(ctx, (NettyConnectMessage) message);
            case CONN_ACK:
            case DISCONNECT:
                return encodeReasonCodeMessage(ctx, message);
            case PUBLISH:
                return encodeSendMessage(ctx, (NettyPublishMessage) message);
            case PUB_ACK:
            case PUB_REC:
            case PUB_REL:
            case PUB_COMP:
                return encodeSendReplyMessage(ctx, message);
            case PING:
            case PONG:
                return encodeMessageWithOnlySingleByteFixedHeader(ctx.alloc(), message);
            default:
                throw new IllegalArgumentException("Unknown message type: " + message.fixedHeader().messageType().value());
        }
    }

    private static ByteBuf encodeConnectMessage(ChannelHandlerContext ctx, NettyConnectMessage message) {
        NettyFixedHeader nettyFixedHeader = message.fixedHeader();
        NettyConnectVariableHeader variableHeader = message.variableHeader();

        // Fixed and variable header
        final ByteBuf variableHeaderBuf = ctx.alloc().buffer();
        try {
            // client id
            String clientId = variableHeader.clientId();
            writeExactUTF8String(variableHeaderBuf, clientId, utf8Bytes(clientId));
            // username
            String userName = variableHeader.username();
            writeEagerUTF8String(variableHeaderBuf, userName);
            // password
            String password = variableHeader.password();
            writeEagerUTF8String(variableHeaderBuf, password);

            int payloadBufferSize = 0;
            int variableHeaderBufferSize = variableHeaderBuf.readableBytes();
            int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
            final int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);
            final ByteBuf variablePartBuf = ctx.alloc().buffer(fixedHeaderBufferSize + variableHeaderBufferSize);

            // Fixed Header Write
            variablePartBuf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
            writeVariableLengthInt(variablePartBuf, variableHeaderBufferSize);
            // Variable Header Write
            variablePartBuf.writeBytes(variableHeaderBuf);
            return variablePartBuf;
        } finally {
            variableHeaderBuf.release();
        }
    }

    private static ByteBuf encodeSendMessage(ChannelHandlerContext ctx, NettyPublishMessage message) {
        NettyFixedHeader nettyFixedHeader = message.fixedHeader();
        NettySendVariableHeader variableHeader = message.variableHeader();
        ByteBuf payload = message.payload();

        // Fixed and variable header
        final ByteBuf variableHeaderBuf = ctx.alloc().buffer();
        try {
            // messageId
            String messageId = variableHeader.messageId();
            writeExactUTF8String(variableHeaderBuf, messageId, utf8Bytes(messageId));
            // messageName
            String messageName = variableHeader.messageName();
            writeExactUTF8String(variableHeaderBuf, messageName, utf8Bytes(messageName));
            // accessToken
            String accessToken = variableHeader.accessToken();
            writeEagerUTF8String(variableHeaderBuf, accessToken);
            // contentType
            String contentType = variableHeader.contentType();
            writeExactUTF8String(variableHeaderBuf, contentType, utf8Bytes(contentType));

            int payloadBufferSize = payload == null ? 0 : payload.readableBytes();
            final int variableHeaderBufferSize = variableHeaderBuf.readableBytes();
            int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
            final int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);
            ByteBuf buf = ctx.alloc().buffer(fixedHeaderBufferSize + variableHeaderBufferSize);

            // Fixed Header Write
            buf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
            writeVariableLengthInt(buf, variableHeaderBufferSize);

            // Variable Header Write
            buf.writeBytes(variableHeaderBuf);

            // Variable Payload Write
            if (payload != null) {
                buf.writeBytes(payload.duplicate());
            }
            return buf;
        } finally {
            variableHeaderBuf.release();
        }
    }

    private static ByteBuf encodeSendReplyMessage(ChannelHandlerContext ctx, NettyMessage message) {
        if (message.variableHeader() instanceof NettySendReplyMessageVariableHeader) {
            NettyFixedHeader nettyFixedHeader = message.fixedHeader();
            NettySendReplyMessageVariableHeader variableHeader = (NettySendReplyMessageVariableHeader) message.variableHeader();
            ByteBuf payload = (ByteBuf) message.payload();

            // Fixed and variable header
            final ByteBuf variableHeaderBuf = ctx.alloc().buffer();
            try {
                // messageId
                String messageId = variableHeader.messageId();
                writeExactUTF8String(variableHeaderBuf, messageId, utf8Bytes(messageId));
                // reasonCode
                variableHeaderBuf.writeInt(variableHeader.reasonCode());
                // reasonMsg
                String reasonMsg = variableHeader.reasonMsg();
                writeExactUTF8String(variableHeaderBuf, reasonMsg, utf8Bytes(reasonMsg));

                int payloadBufferSize = payload == null ? 0 : payload.readableBytes();
                final int variableHeaderBufferSize = variableHeaderBuf.readableBytes();
                int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
                final int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);
                ByteBuf buf = ctx.alloc().buffer(fixedHeaderBufferSize + variableHeaderBufferSize);

                // Fixed Header Write
                buf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
                writeVariableLengthInt(buf, variableHeaderBufferSize);

                // Variable Header Write
                buf.writeBytes(variableHeaderBuf);

                // Variable Payload Write
                if (payload != null) {
                    buf.writeBytes(ProtostuffSerializeUtils.serialize(payload.duplicate()));
                }
                return buf;
            } finally {
                variableHeaderBuf.release();
            }
        } else {
            return encodeMessageWithOnlySingleByteFixedHeaderAndMessageId(ctx.alloc(), message);
        }
    }

    private static ByteBuf encodeMessageWithOnlySingleByteFixedHeaderAndMessageId(ByteBufAllocator byteBufAllocator, NettyMessage message) {
        NettyFixedHeader nettyFixedHeader = message.fixedHeader();
        NettyMessageIdVariableHeader variableHeader = (NettyMessageIdVariableHeader) message.variableHeader();
        // Fixed and variable header
        final ByteBuf variableHeaderBuf = byteBufAllocator.buffer();
        try {
            // messageId
            String messageId = variableHeader.messageId();
            writeExactUTF8String(variableHeaderBuf, messageId, utf8Bytes(messageId));

            int payloadBufferSize = 0;
            int variableHeaderBufferSize = variableHeaderBuf.readableBytes();
            int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
            int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);
            ByteBuf buf = byteBufAllocator.buffer(fixedHeaderBufferSize + variableHeaderBufferSize);

            // Fixed Header Write
            buf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
            writeVariableLengthInt(buf, variableHeaderBufferSize);

            // Variable Header Write
            buf.writeBytes(variableHeaderBuf);
            return variableHeaderBuf;
        } finally {
            variableHeaderBuf.release();
        }
    }

    private static ByteBuf encodeReasonCodeMessage(ChannelHandlerContext ctx, NettyMessage message) {
        if (message.variableHeader() instanceof NettyReasonCodeVariableHeader) {
            NettyFixedHeader nettyFixedHeader = message.fixedHeader();
            NettyReasonCodeVariableHeader variableHeader = (NettyReasonCodeVariableHeader) message.variableHeader();

            // Fixed and variable header
            final ByteBuf reasonMsgBuf = encodeString(ctx.alloc(), variableHeader.reasonMsg());
            try {
                final int reasonMsgBufferSize = reasonMsgBuf.readableBytes();
                final int reasonCodeBufferSize = 4;
                int payloadBufferSize = 0;
                final int variableHeaderBufferSize = reasonCodeBufferSize + reasonMsgBufferSize;
                int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
                final int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);
                ByteBuf buf = ctx.alloc().buffer(fixedHeaderBufferSize + variableHeaderBufferSize);

                // Fixed Header Write
                buf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
                writeVariableLengthInt(buf, variableHeaderBufferSize);

                // Variable Header Write
                buf.writeInt(variableHeader.reasonCode());
                buf.writeBytes(reasonMsgBuf);
                return buf;
            } finally {
                reasonMsgBuf.release();
            }
        } else {
            return encodeMessageWithOnlySingleByteFixedHeader(ctx.alloc(), message);
        }
    }

    private static ByteBuf encodeMessageWithOnlySingleByteFixedHeader(ByteBufAllocator byteBufAllocator, NettyMessage message) {
        NettyFixedHeader nettyFixedHeader = message.fixedHeader();
        ByteBuf buf = byteBufAllocator.buffer(2);
        buf.writeByte(getFixedHeaderByte1(nettyFixedHeader));
        buf.writeByte(0);
        return buf;
    }

    private static ByteBuf encodeString(ByteBufAllocator byteBufAllocator, String message) {
        ByteBuf stringBuf = byteBufAllocator.buffer();
        NettyCodecUtil.writeEagerUTF8String(stringBuf, message);
        return stringBuf;
    }
}
